package security

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"os"
	"time"

	"github.com/golang-jwt/jwt/v4"
)

// JWTManager handles JWT token operations securely
type JWTManager struct {
	secretKey     []byte
	issuer        string
	tokenExpiry   time.Duration
	refreshExpiry time.Duration
}

// Claims represents JWT claims with security enhancements
type Claims struct {
	UserID   int64  `json:"user_id"`
	Username string `json:"username"`
	Role     string `json:"role"`
	TokenID  string `json:"token_id"` // For token revocation
	jwt.RegisteredClaims
}

// TokenPair represents access and refresh tokens
type TokenPair struct {
	AccessToken  string    `json:"access_token"`
	RefreshToken string    `json:"refresh_token"`
	ExpiresAt    time.Time `json:"expires_at"`
	TokenType    string    `json:"token_type"`
}

// NewJWTManager creates a secure JWT manager
func NewJWTManager() (*JWTManager, error) {
	// Use environment variable for JWT secret, never hardcode
	secretKey := os.Getenv("JWT_SECRET_KEY")
	if secretKey == "" {
		// Generate a secure random key if not provided (for development only)
		key := make([]byte, 32)
		if _, err := rand.Read(key); err != nil {
			return nil, err
		}
		secretKey = hex.EncodeToString(key)
		// In production, this should cause an error
		// return nil, errors.New("JWT_SECRET_KEY environment variable not set")
	}

	return &JWTManager{
		secretKey:     []byte(secretKey),
		issuer:        "gin-framework",
		tokenExpiry:   15 * time.Minute, // Short-lived access tokens
		refreshExpiry: 24 * time.Hour,   // Longer-lived refresh tokens
	}, nil
}

// GenerateTokenPair creates both access and refresh tokens
func (m *JWTManager) GenerateTokenPair(user *SecureUser) (*TokenPair, error) {
	// Generate unique token ID for revocation tracking
	tokenID, err := m.generateTokenID()
	if err != nil {
		return nil, err
	}

	now := time.Now()
	expiresAt := now.Add(m.tokenExpiry)

	// Create access token claims
	claims := &Claims{
		UserID:   user.ID,
		Username: user.Username,
		Role:     user.Role,
		TokenID:  tokenID,
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    m.issuer,
			Subject:   user.Username,
			Audience:  []string{"gin-framework"},
			ExpiresAt: jwt.NewNumericDate(expiresAt),
			NotBefore: jwt.NewNumericDate(now),
			IssuedAt:  jwt.NewNumericDate(now),
			ID:        tokenID,
		},
	}

	// Create and sign access token
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	accessToken, err := token.SignedString(m.secretKey)
	if err != nil {
		return nil, err
	}

	// Create refresh token with longer expiry
	refreshTokenID, err := m.generateTokenID()
	if err != nil {
		return nil, err
	}

	refreshClaims := &Claims{
		UserID:   user.ID,
		Username: user.Username,
		Role:     user.Role,
		TokenID:  refreshTokenID,
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    m.issuer,
			Subject:   user.Username,
			Audience:  []string{"gin-framework-refresh"},
			ExpiresAt: jwt.NewNumericDate(now.Add(m.refreshExpiry)),
			NotBefore: jwt.NewNumericDate(now),
			IssuedAt:  jwt.NewNumericDate(now),
			ID:        refreshTokenID,
		},
	}

	refreshToken := jwt.NewWithClaims(jwt.SigningMethodHS256, refreshClaims)
	refreshTokenString, err := refreshToken.SignedString(m.secretKey)
	if err != nil {
		return nil, err
	}

	return &TokenPair{
		AccessToken:  accessToken,
		RefreshToken: refreshTokenString,
		ExpiresAt:    expiresAt,
		TokenType:    "Bearer",
	}, nil
}

// ValidateToken validates and parses a JWT token
func (m *JWTManager) ValidateToken(tokenString string) (*Claims, error) {
	// Parse token with claims
	token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
		// Ensure signing method is what we expect
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, errors.New("unexpected signing method")
		}
		return m.secretKey, nil
	})

	if err != nil {
		return nil, err
	}

	// Check if token is valid
	if !token.Valid {
		return nil, errors.New("invalid token")
	}

	// Extract claims
	claims, ok := token.Claims.(*Claims)
	if !ok {
		return nil, errors.New("invalid token claims")
	}

	// Additional validation
	if claims.Issuer != m.issuer {
		return nil, errors.New("invalid token issuer")
	}

	// Check if token is expired (additional check)
	if time.Now().After(claims.ExpiresAt.Time) {
		return nil, errors.New("token expired")
	}

	return claims, nil
}

// RefreshToken creates a new access token using a valid refresh token
func (m *JWTManager) RefreshToken(refreshTokenString string, userRepo UserRepository) (*TokenPair, error) {
	// Validate refresh token
	claims, err := m.ValidateToken(refreshTokenString)
	if err != nil {
		return nil, err
	}

	// Check if this is actually a refresh token
	if len(claims.Audience) == 0 || claims.Audience[0] != "gin-framework-refresh" {
		return nil, errors.New("invalid refresh token")
	}

	// Get current user data (for role changes, active status, etc.)
	user, err := userRepo.GetUserByUsername(claims.Username)
	if err != nil {
		return nil, err
	}

	if !user.IsActive {
		return nil, errors.New("user account is disabled")
	}

	// Generate new token pair
	return m.GenerateTokenPair(user)
}

// generateTokenID creates a unique token identifier
func (m *JWTManager) generateTokenID() (string, error) {
	bytes := make([]byte, 16)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}
	return hex.EncodeToString(bytes), nil
}

// RevokeToken marks a token as revoked (requires Redis/database implementation)
func (m *JWTManager) RevokeToken(tokenID string) error {
	// Implementation would store revoked token IDs in Redis/database
	// This is a placeholder for the actual implementation
	// revokedTokens.Set(tokenID, time.Now(), remaining_ttl)
	return nil
}

// IsTokenRevoked checks if a token has been revoked
func (m *JWTManager) IsTokenRevoked(tokenID string) bool {
	// Implementation would check revoked token IDs in Redis/database
	// This is a placeholder for the actual implementation
	// return revokedTokens.Exists(tokenID)
	return false
}